import os

from Katana import Nodes3DAPI, RenderingAPI
from Nodes3DAPI import RenderNodeUtil
from Nodes3DAPI import RenderingUtil
from Nodes3DAPI import RenderNodeObject

#
# ArnoldRenderNodeObject - Handle details for arnold render
#
class ArnoldRenderNodeObject(object):
    # Get the render settings and filter text
    def __get_render_details(self):
        # Note: just in case it comes up again, we cannot use the producer or
        # render settings cached in self.__renderNodeInfo.  It is a geometry producer,
        # not a render producer.
        self.__producer = Nodes3DAPI.GetRenderProducer(self.__node, self.__time, True,
                    applyRenderSelectedMask=False)

        # render settings
        self.__renderSettings = RenderNodeUtil.GetRenderSettings(self.__producer,
            cameraErrorIsFatal=True)

        # We happen to know the display/data windows of the render here.        
        self.__displayWindow = self.__renderSettings['display window']
        self.__dataWindow = self.__renderSettings['data window']
        self.__finalDisplayWindow = self.__renderSettings['final display window']
        self.__finalDataWindow = self.__renderSettings['final data window']

        # filter text
        self.__filterText = Nodes3DAPI.GetProducerText(self.__producer)

        # Checkpoint data is accumulated over each of the outputs by scanning for
        # files from previous renders
        self.__checkpointData = {'outputs':{}}
        self.__checkpointEnable = False
        if 'KATANA_BATCH' in os.environ:
            checkpointRenderAttr = self.__producer.getDelimitedLocalAttribute(
                'arnoldGlobalStatements.checkpointRender')
            if checkpointRenderAttr and bool(checkpointRenderAttr.getData()[0]):
                self.__checkpointEnable = True

    # Generate .py and .xml filenames so they can be used inside arnold
    def __init_tmp_filenames(self):
        self.__filterProgramFilename = RenderingUtil.GetTmpLocalFile('.py', 'renderProducerText')
        self.__productStateTableFilename = RenderingUtil.GetTmpLocalFile('.xml', 'productStateTable')
        self.__outputInfoFilename = RenderingUtil.GetTmpLocalFile('.xml', 'outputInfo')
        self.__renderFinishedFilename = RenderingUtil.GetTmpLocalFile('.txt', 'renderFinished')
        self.__cleanupFiles.append(self.__filterProgramFilename)
        self.__cleanupFiles.append(self.__productStateTableFilename)
        self.__cleanupFiles.append(self.__outputInfoFilename)
        self.__cleanupFiles.append(self.__renderFinishedFilename)

    # Write out the producer to a .py file so it can be used inside the renderer
    def __write_producer_file(self):
        if not self.__filterProgramFilename: return

        filterProgramFile = file(self.__filterProgramFilename, "w")
        filterProgramFile.write(self.__filterText)
        filterProgramFile.close()
        RenderingUtil.SetModeFlags(self.__filterProgramFilename)

    # Write out the product state table to a .xml file so it can be used inside the renderer
    def __write_product_state_table(self):
        if not self.__productStateTableFilename: return

        RenderNodeObject._ExportProductTables(self.__productStateTableFilename)
        RenderingUtil.SetModeFlags(self.__productStateTableFilename)

    def __write_output_info_file(self):
        if not self.__outputInfoFilename: return

        outputInfoFile = file(self.__outputInfoFilename, "w")
        outputInfoFile.write('<render node="%s" renderFinishedFilename="%s">\n' % (
            self.__node.getName(), self.__renderFinishedFilename))
        for outputInfo in self.__renderOutputs:
            outputInfoFile.write('  <output name="%s" file="%s" channel="%s"/>\n' % outputInfo)
        if self.__checkpointData:
            outputInfoFile.write('  <checkpoint tileXSize="%d" tileYSize="%d"' % \
                self.__checkpointData.get('tileSize', (0, 0)))
            outputInfoFile.write(' numXTiles="%d" numYTiles="%d">\n' % \
                self.__checkpointData.get('numTiles', (0, 0)))

            for name, output in self.__checkpointData['outputs'].iteritems():
                outputInfoFile.write('    <output name="%s"' % name)
                outputInfoFile.write(' checkpointFile="%s"' % output['checkpointFile'])
                outputInfoFile.write(' validTiles="')
                
                validTileList = list(output['validTiles'])
                # sort by y, then by x
                validTileList.sort(lambda a, b: a[1]-b[1] or a[0]-b[0])
                for validTile in validTileList:
                    outputInfoFile.write('%d,%d;' % validTile)

                outputInfoFile.write('"/>\n')
            outputInfoFile.write('  </checkpoint>\n')
        outputInfoFile.write('</render>\n')
        outputInfoFile.close()
        RenderingUtil.SetModeFlags(self.__outputInfoFilename)

    # Build a arnold command line.
    def __build_arnold_commandline(self):
        envLine = ''
        katanaPythonPath = os.environ.get('KATANA_PYTHON_LIBS', '')
        pythonPath = os.environ.get('PYTHONPATH', '')
        if katanaPythonPath != '':
            if pythonPath != '':
                pythonPath = katanaPythonPath+':'+pythonPath
            else:
                pythonPath = katanaPythonPath
        if pythonPath != '':
            envLine += 'PYTHONPATH='+pythonPath+' '

        katanaSpImportPackages = os.environ.get('KATANA_SPIMPORT_PACKAGES', '')
        spImportPackages = os.environ.get('SPIMPORT_PACKAGES', '')
        if spImportPackages != '':
            spImportPackages = katanaSpImportPackages+':'+spImportPackages
        else:
            spImportPackages = katanaSpImportPackages
        envLine += 'SPIMPORT_PACKAGES='+spImportPackages+' '
        envLine += 'SPIMPORT_NOWARNING=1 '
        
        libPath = os.environ.get('LD_LIBRARY_PATH','')
        arnoldLibPath = os.environ.get('ARNOLD_LD_LIBRARY_PATH','')
        envLine += 'LD_LIBRARY_PATH=%s:%s:%s ' % (arnoldLibPath, libPath, RenderingAPI.RenderPlugins.GetRendererPluginBinDir("arnold"))
        
        envLine += RenderNodeObject._GetAssetEnvString()
        
        if envLine:
            envLine = 'env '+envLine
        
        
        # Get the arnold executable
        pluginPath = RenderingAPI.RenderPlugins.GetRendererPluginBinDir("arnold")
        arnoldRender = os.path.join(pluginPath, "arnold_render")

        cmd = envLine
        cmd += arnoldRender

        cmd += " --producerText %s" % (self.__filterProgramFilename)
        cmd += " --productStateTable %s" % (self.__productStateTableFilename)
        cmd += " --outputInfo %s" % (self.__outputInfoFilename)
        cmd += " --width %s" % (str(self.__dataWindow[2]-self.__dataWindow[0]))
        cmd += " --height %s" % (str(self.__dataWindow[3]-self.__dataWindow[1]))
        cmd += " <CROPWINDOW>"
        cameraSettings = self.__renderSettings['camera settings']
        cmd += " --screenWindow %f,%f,%f,%f" % (cameraSettings['screen window'][0],
                                                cameraSettings['screen window'][2],
                                                cameraSettings['screen window'][1],
                                                cameraSettings['screen window'][3])
        cmd += " --frameTime %s" % (str(self.__time))
        cmd += " --batchRender"

        logCommand = os.environ.get('KATANA_BASE', 'error') + "/MAIN/bin/exec_and_log"
        logCommand += " --package arnold"
        logCommand += " --version %s" % (os.environ.get('ARNOLD_VERSION', 'unknown'))
        logCommand += " --apppath %s" % (arnoldRender)
        logCommand += " -- "

        self.__commandLines.append(logCommand + cmd)
        
        if os.environ.get('ARNOLD_VERSION', '').startswith('2.'):
            raise RuntimeError, 'This version of katana requires Arnold 3 or higher.'
        
        
    # Process the defined outputs here, building up the commands and such.
    def __process_outputs(self):
        totalEnabled = 0

        for i in xrange(self.__renderNodeInfo.getNumberOfOutputs()):
            outputInfo = self.__renderNodeInfo.getOutputInfoByIndex(i, self.__forceLocal,
                makeVersionTemplate=False)
            outputName = outputInfo['name']
            outputType = outputInfo['type']

            # add enabled status to cacheID
            self.__cacheID += outputName+'('+str(outputInfo['enabled'])+')'

            if not outputInfo['enabled']:
                self.__add_dummy_output()
                continue

            totalEnabled += 1

            # Do something with the type...
            if outputType == 'raw':
                self.__process_raw_output(outputInfo)
            elif outputType == 'color':
                self.__process_color_output(outputInfo)
            elif outputType == 'shadow':
                self.__process_shadow_output(outputInfo)
            elif outputType == 'script':
                self.__process_script_output(outputInfo)
            elif outputType == 'prescript':
                self.__process_prescript_output(outputInfo)
            elif outputType == 'merge':
                self.__process_merge_output(outputInfo)
            else:
                errmsg = "Error processing output '%s' for Arnold render: " % (outputName)
                errmsg += " unsupported output type '%s'." %  (outputType)
                raise RuntimeError, errmsg

        if totalEnabled == 0:
            raise RuntimeError, "No enabled outputs to render."

    def __add_dummy_output(self):
        self.__imageFiles.append('')

    def __process_raw_output(self, outputInfo):
        fileExtension = outputInfo['fileExtension']
        localOutputFile = outputInfo['tempRenderLocation']
        if not localOutputFile:
            localOutputFile = RenderingUtil.GetTmpLocalFile('.'+fileExtension, 'render')
        self.__cleanupFiles.append(localOutputFile)
        
        outputChannel = outputInfo['channel']

        self.__renderOutputs.append((outputInfo['name'], localOutputFile, outputChannel))

        rawHasOutput = outputInfo['rawHasOutput']
        if rawHasOutput:
            if self.__renderNodeInfo.isTileRender():
                tileOutputFile = self.__renderNodeInfo.computeTileFileName(outputInfo)
                self.__convertCommands.append('cp -f %(localOutputFile)s %(tileOutputFile)s' % locals())
                self.__imageFiles.append(tileOutputFile)
            else:
                finalOutputFile = outputInfo['outputFile']
                self.__convertCommands.append('cp -f %(localOutputFile)s %(finalOutputFile)s' % locals())
                self.__imageFiles.append(finalOutputFile)
        else:
            self.__imageFiles.append(RenderingUtil.GetTmpLocalFile('.doNotLoadMe', 'placeHolder'))

    def __process_color_output(self, outputInfo):
        
        localOutputFile = outputInfo.get('localFile', None)
        if not localOutputFile:
            localOutputFile = RenderingUtil.GetTmpLocalFile('.exr', 'render')
        self.__cleanupFiles.append(localOutputFile)

        outputChannel = outputInfo['channel']

        # sub in info from render camera, if present
        additionalExrAttrs = outputInfo['additionalExrAttrs'].copy()

        cameraName = outputInfo.get('customCameraName', None)
        if cameraName:
            for k in additionalExrAttrs.keys():
                if k.startswith('spi:camera:'):
                    del additionalExrAttrs[k]
            try:
                cameraProducer = self.__producer.getProducerByPath(cameraName)
                cameraInfoDict = RenderNodeUtil.GetCameraInfoDict(cameraProducer)

                for k,v in cameraInfoDict.iteritems():
                    additionalExrAttrs['spi:camera:'+k] = str(v)
            except:
                pass

        self.__renderOutputs.append((outputInfo['name'], localOutputFile, outputChannel))

        finalConvertCommand = None

        if self.__renderNodeInfo.isTileRender():
            tileOutputFile = self.__renderNodeInfo.computeTileFileName(outputInfo)
            checkpointOutputFile = self.__compute_checkpoint_output_file(tileOutputFile)
            finalConvertCommand = 'cp -f %(localOutputFile)s %(tileOutputFile)s' % locals()
            self.__imageFiles.append(tileOutputFile)
        else:
            finalOutputFile = outputInfo['outputFile']
            checkpointOutputFile = self.__compute_checkpoint_output_file(finalOutputFile)
            self.__convertFiles.append((localOutputFile, finalOutputFile,
                outputInfo['convertSettings'], outputInfo['clampOutput'],
                outputInfo['colorConvert'], additionalExrAttrs, outputInfo['computeStats']))
            self.__imageFiles.append(finalOutputFile)

        if self.__accum_checkpoint_renders(outputInfo['name'], checkpointOutputFile):
            tmpCheckpointFile = RenderingUtil.GetTmpLocalFile('.exr', 'checkpoint')
            checkpointCmd = os.environ.get('KATANA_BASE', 'error') + '/MAIN/bin/checkpointstitch -cleanup '
            checkpointCmd += localOutputFile + ' ' + checkpointOutputFile + ' '
            checkpointCmd += tmpCheckpointFile
            self.__convertCommands.append(checkpointCmd)
            self.__convertCommands.append('mv %(tmpCheckpointFile)s %(localOutputFile)s' % locals())

        if finalConvertCommand:
            self.__convertCommands.append(finalConvertCommand)


    def __process_shadow_output(self, outputInfo):
        shadowType = outputInfo['shadowType']

        if shadowType not in ('deepshad', 'vshd'):
            errmsg = "Error processing output '%s' for Arnold render: " % (outputInfo['name'])
            errmsg += " unsupported shadow output type '%s'." %  (shadowType)
            raise RuntimeError, errmsg

        # Process deepshad/vshd.
        finalOutputFile = outputInfo['outputFile']
        outputChannel = outputInfo['channel']

        localOutputFile = outputInfo['tempRenderLocation']
        if not localOutputFile:
            localOutputFile = RenderingUtil.GetTmpLocalFile('.%s'%shadowType, 'render')
        self.__cleanupFiles.append(localOutputFile)
        self.__renderOutputs.append((outputInfo['name'], localOutputFile, outputChannel))

        finalConvertCommand = None
        finalOverscanCommand = None
        if self.__renderNodeInfo.isTileRender():
            tileOutputFile = self.__renderNodeInfo.computeTileFileName(outputInfo)
            checkpointOutputFile = self.__compute_checkpoint_output_file(tileOutputFile)
            finalConvertCommand = 'cp -f %(localOutputFile)s %(tileOutputFile)s' % locals()
            self.__imageFiles.append(tileOutputFile)
        else:
            if shadowType == 'deepshad':
                finalConvertCommand = 'cp -f %(localOutputFile)s %(finalOutputFile)s' % locals()
                self.__imageFiles.append(finalOutputFile)
                checkpointOutputFile = self.__compute_checkpoint_output_file(finalOutputFile)
            elif shadowType == 'vshd':
                intermediateFile = RenderingUtil.GetTmpLocalFile('.%s'%shadowType, 'render')
                finalConvertCommand = 'cp -f %(localOutputFile)s %(intermediateFile)s' % locals()

                if self.__renderSettings.get('overscan',(0,0,0,0)) != (0,0,0,0):
                    header = {}
                    dataWindow = self.__renderSettings['final data window']
                    displayWindow = self.__renderSettings['final display window']
                    shadowconvertVersion = os.environ.get('SHADOWCONVERT_VERSION', '')
                    if shadowconvertVersion:
                        shadowconvertVersion = '-appver ' + shadowconvertVersion + ' '
                    finalOverscanCommand = ('shadowconvert %(shadowconvertVersion)s' \
                                         +  '%(intermediateFile)s %(finalOutputFile)s ' \
                                         +  '-dataOffset %(x)i %(y)i ' \
                                         +  '-displayWindow %(dx0)i %(dy0)i %(dx1)i %(dy1)i ' \
                                         +  '-fixWorldToNDC '\
                                         +  '-uncroppedDataWindow %(udx0)i %(udy0)i %(udx1)i %(udy1)i') % {
                                            'shadowconvertVersion' : shadowconvertVersion,
                                            'finalOutputFile' : finalOutputFile,
                                            'intermediateFile': intermediateFile,
                                            'x'               : dataWindow[0],
                                            'y'               : displayWindow[3]-dataWindow[3], # inverted y
                                            'dx0'             : displayWindow[0],
                                            'dy0'             : displayWindow[1],
                                            'dx1'             : displayWindow[2],
                                            'dy1'             : displayWindow[3],
                                            'udx0'            : dataWindow[0],
                                            'udy0'            : displayWindow[3]-dataWindow[3], # inverted y
                                            'udx1'            : dataWindow[2],
                                            'udy1'            : displayWindow[3]-dataWindow[1], # inverted y
                                         }
                else:
                    finalOverscanCommand = 'mv %(intermediateFile)s %(finalOutputFile)s' % locals()

                self.__imageFiles.append(finalOutputFile)
                checkpointOutputFile = self.__compute_checkpoint_output_file(finalOutputFile)

        if self.__accum_checkpoint_renders(outputInfo['name'], checkpointOutputFile):
            tmpCheckpointFile = RenderingUtil.GetTmpLocalFile('.%s'%shadowType, 'checkpoint')
            checkpointCmd = os.environ.get('KATANA_BASE', 'error') + '/MAIN/bin/checkpointstitch -cleanup '
            checkpointCmd += localOutputFile + ' ' + checkpointOutputFile + ' '
            checkpointCmd += tmpCheckpointFile
            self.__convertCommands.append(checkpointCmd)
            self.__convertCommands.append('mv %(tmpCheckpointFile)s %(localOutputFile)s' % locals())


        if finalConvertCommand:
            self.__convertCommands.append(finalConvertCommand)
        if finalOverscanCommand:
            self.__convertCommands.append(finalOverscanCommand)

    def __process_script_output(self, outputInfo):
        # for tile render, don't do anything.  save it for the stitch pass.
        if self.__renderNodeInfo.isTileRender():
            self.__imageFiles.append('')
        else:
            postCommand = outputInfo['scriptCommand']
            if postCommand != '':
                self.__postCommands.extend(postCommand.split(';'))
            if outputInfo['scriptHasOutput']:
                self.__imageFiles.append(outputInfo['outputFile'])
            else:
                self.__imageFiles.append(outputInfo['passThroughInput'])

    def __process_merge_output(self, outputInfo):
        # don't attempt to load merged image (wait for multi-layer support in 2D system)
        self.__imageFiles.append('')

    def __process_prescript_output(self, outputInfo):
        # for tile render, don't do anything.  (TODO: need prescripts to run during tile setup)
        if not self.__renderNodeInfo.isTileRender():
            preCommand = outputInfo['scriptCommand']
            if preCommand != '':
                self.__preCommands.extend(
                    filter(bool, [x.strip() for x in preCommand.split(';')]))
        self.__imageFiles.append(outputInfo['outputFile'])

    def __compute_checkpoint_output_file(self, outputFile):
        fileDir = os.path.dirname(outputFile)
        fileBase = os.path.basename(outputFile)
        baseSplit = fileBase.split('.')

        tmpDir = os.environ.get('KATANA_TMPDIR', '/tmp')
        show = os.environ.get('SHOW', None)
        shot = os.environ.get('SHOT', None)
        if show and shot and fileDir.startswith(tmpDir):
            fileDir = os.getenv('KATANA_SHOTS')+'/%s/%s/pix/checkpoint/%s' % (show, shot, baseSplit[0])
            return fileDir + '/' + fileBase + '.CHECKPOINT'
        else:
            return fileDir + '/.checkpoint/' + fileBase + '.CHECKPOINT'

    def __accum_checkpoint_renders(self, outputName, checkpointFile):
        if not self.__checkpointEnable: return False
        if not checkpointFile.endswith('.CHECKPOINT'): return False

        allTiles = None
        validTiles = set()
        needStitch = False

        # walk through all checkpoint files, find valid tiles
        # (union of valid tiles from all checkpoint images for current image)
        i = 1
        while True:
            findFile = checkpointFile.rsplit('.', 1)[0] + '.checkpoint.%d'%i
            i += 1

            if not os.path.exists(findFile): break
            if not PyOpenEXR.IsExrFile(findFile): continue

            try:
                tileDesc = PyOpenEXR.GetExrTileDescription(findFile)
            except Exception, e:
                log.warning('Ignoring checkpoint file: %s' % e)
                continue

            tileSize = (tileDesc['tileXSize'], tileDesc['tileYSize'])
            numTiles = (tileDesc['numXTiles'], tileDesc['numYTiles'])

            if allTiles is None:
                allTiles = set()
                for y in xrange(numTiles[1]):
                    for x in xrange(numTiles[0]):
                        allTiles.add((x, y))

            if 'tileSize' in self.__checkpointData:
                if tileSize != self.__checkpointData['tileSize']:
                    raise RuntimeError, 'Differing tile size in checkpoint files'
            else:
                self.__checkpointData['tileSize'] = tileSize

            if 'numTiles' in self.__checkpointData:
                if numTiles != self.__checkpointData['numTiles']:
                    raise RuntimeError, 'Differing number of tiles in checkpoint files'
            else:
                self.__checkpointData['numTiles'] = numTiles

            # For all checkpoint tiles related to a single image, accumulate
            # valid tiles that exist in any file.
            invalidTiles = set(tileDesc['invalidTiles'])
            validTiles = validTiles.union(allTiles.difference(invalidTiles))

            # At least one valid checkpoint file... stitch is needed.
            needStitch = True

        self.__checkpointData['outputs'][outputName] = {
            'validTiles': validTiles,
            'checkpointFile': checkpointFile}

        return needStitch


    # This takes care of subsituting the CropWindow into the commandline.
    def __modify_cropwindow_in_commandline(self, rect):
        cropWindow = RenderingUtil.CalcCropWindow(rect, self.__dataWindow,
            self.__renderSettings['crop window'])

        cropWindowString = ''
        if cropWindow != (0.0, 1.0, 0.0, 1.0):
            # unfortunate ugly to workaround currently unignorable
            # pychecker warning.
            cropWindowString = ' --cropWindow %f,%f,%f,%f'.__mod__(cropWindow)

        self.__commandLines = [x.replace('<CROPWINDOW>', cropWindowString) for x in self.__commandLines]

    def _real_init(self, node, time, hotRender):
        self.__node = node
        self.__time = time

        # If we're bypassed, pretend we're hotrendering...
        # We won't actually run any commands or anything, but
        # we need to know where we'll write to the shot tree
        # (so we can load that file).
        # I'm sorry, but I think this subtle interaction is unavoidable.
        self.__forceLocal = not (hotRender or node.isBypassed())

        self.__preCommands = [] # Commands to run before anything else
        self.__commandLines = []
        self.__postCommands = []  # Commands to run after all convert commands have run
        self.__convertCommands = [] # commands to run after the renderer finishes (one per pass)
        self.__imageFiles = []  # final images to keep around
        self.__convertFiles = [] # list of (from, to, outputArgs) tuples for IMGLIB convert
        self.__cleanupFiles = []  # images to delete after everything is done

        self.__renderNodeInfo = RenderNodeUtil.GetRenderNodeInfo(self.__node, self.__time)
        self.__cacheID = self.__renderNodeInfo.getCacheID()

        self.__renderOutputs = []

        self.__get_render_details()
        self.__process_outputs()

        self.__init_tmp_filenames()
        
        # Build all needed structures for arnold running....
        self.__build_arnold_commandline()
        
        for cmd in self.__convertCommands:
            self.__commandLines.append(cmd)
        del self.__convertCommands


    def __init__(self, node, time, hotRender):
        # For development purposes, we wrap this in a try/except.
        # _real_init is used to avoid having to indent+unindent all the code.
        try:
            self._real_init(node, time, hotRender)
        except:
            raise

    def __del__(self):
        pass

    # Write out the final producer text and RIBStream for rendering.
    def finalizeSetup(self, rect):
        self.__write_producer_file()
        self.__write_product_state_table()
        self.__write_output_info_file()
        self.__modify_cropwindow_in_commandline(rect)

    def getPreCommands(self): return self.__preCommands
    def getCommandLines(self): return self.__commandLines
    def getPostCommands(self): return self.__postCommands
    def getImageFiles(self): return self.__imageFiles
    def getCleanupFiles(self): return self.__cleanupFiles
    def getConvertFiles(self): return self.__convertFiles
    def getRenderFinishedFilename(self): return self.__renderFinishedFilename
    def getDataWindow(self): return self.__finalDataWindow
    def getDisplayWindow(self): return self.__finalDisplayWindow
    def getCacheID(self): return self.__cacheID
    def getDisplayStatements(self): return ""           

